import pyE17
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import scipy.fftpack as fft
import numpy as np
def fgrid(sh, psize=None):
    """\
    ``q0,q1,... = fgrid(sh)``
    returns Fourier-space coordinates for a N-dimensional array of shape sh (pixel units)

    ``q0,q1,... = fgrid(sh,psize)``
    gives the coordinates according to the given pixel size psize.
    
    Parameters
    ----------
    sh : the shape of the N-dimensional array
    
    psize : pixel size in each dimensions
    
    """
    if psize is None:
        return fft.ifftshift(np.indices(sh).astype(float) - np.reshape(np.array(sh) // 2, (len(sh),) + len(sh) * (1,)), range(1, len(sh) + 1))
    else:
        psize = np.asarray(psize)
        if psize.size == 1:
            psize = psize * np.ones((len(sh),))
        psize = np.asarray(psize).reshape((len(sh),) + len(sh) * (1,))
        return fft.ifftshift(np.indices(sh).astype(float) - np.reshape(np.array(sh) // 2, (len(sh),) + len(sh) * (1,)), range(1, len(sh) + 1)) * psize

def fvec2(sh, psize=None):
    """\
    q2 = fvec2(sh, psize):
    squared norm of reciprocal space coordinates, with pixel size psize.
    """
    return np.sum(fgrid(sh, psize) ** 2, axis=0)
def free_nf(w, l, z, pixsize=1.):
    """\
    Free-space propagation (near field) of the wavefield of a distance z.
    l is the wavelength. 
    """
#    if w.ndim != 2:
#        raise RunTimeError("A 2-dimensional wave front 'w' was expected")

    sh = w.shape

    #if sh[0] != sh[1]:
    #    raise RunTimeError("Only implemented for square arrays...")

    # Convert to pixel units.
    z = z / pixsize
    l = l / pixsize

    # Evaluate if aliasing could be a problem
    if min(sh) / np.sqrt(2.) < z * l:
        print "Warning: z > N/(sqrt(2)*lamda) = %.6g: this calculation could fail." % (min(sh) / (l * np.sqrt(2.))) 
        print "(consider padding your array, or try a far field method)"  

    #q2 = U.fvec2(sh)
    #return np.fft.ifftn(np.fft.fftn(w) * np.exp(2j * np.pi * (z / l) * (np.sqrt(1 - q2*(l/sh[0])**2) - 1) ) )

    q2 = fvec2(sh, psize=(1. / sh[0], 1. / sh[1]))
    return np.fft.ifftn(np.fft.fftn(w) * np.exp(2j * np.pi * (z / l) * (np.sqrt(1 - q2 * l ** 2) - 1)))

p = pyE17.utils.imload('pinhole.png')
probe = p# free_nf(p, 2e-10, 0.7e-3, 20e-9)
p = pyE17.utils.imsave(probe, 'pinhole_fresnel.png')
plt.imshow(probe.real, interpolation='bilinear', cmap=cm.gray)
plt.imshow(probe.imag, interpolation='bilinear', cmap=cm.gray)
#plt.show()

pyE17.io.h5write('probe_0.h5', probe=probe) 
